import org.bouncycastle.asn1.x500.X500Name;
import org.bouncycastle.cert.X509CertificateHolder;
import org.bouncycastle.cert.jcajce.JcaX509CertificateConverter;
import org.bouncycastle.cert.ocsp.*;
import org.bouncycastle.cert.ocsp.jcajce.JcaCertificateID;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.operator.ContentSigner;
import org.bouncycastle.operator.jcajce.JcaContentSignerBuilder;
import org.bouncycastle.operator.jcajce.JcaDigestCalculatorProviderBuilder;
import org.bouncycastle.openssl.PEMParser;
import org.bouncycastle.openssl.jcajce.JcaPEMKeyConverter;
import org.bouncycastle.asn1.gm.GMObjectIdentifiers;
import org.bouncycastle.asn1.x509.AlgorithmIdentifier;
import org.bouncycastle.cert.ocsp.CertificateID;
import org.bouncycastle.cert.ocsp.OCSPException;
import org.bouncycastle.operator.DigestCalculator;
import org.bouncycastle.operator.jcajce.JcaDigestCalculatorProviderBuilder;

import java.io.*;
import java.math.BigInteger;
import java.security.PrivateKey;
import java.security.Security;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;
import java.util.Date;

public class SimpleOCSPServer_SM2 {

    static {
        // 注册Bouncy Castle提供程序以支持国密算法
        Security.addProvider(new BouncyCastleProvider());
    }

    public static void main(String[] args) throws Exception {
        // 1. 加载证书和私钥
        X509Certificate caCert = loadCert("CA.crt");
        PrivateKey caPrivateKey = loadKey("CA.key");
        X509Certificate entityCert = loadCert("b.crt");
        
        if (caCert == null) {
            System.err.println("无法加载CA证书");
            return;
        }
        if (caPrivateKey == null) {
            System.err.println("无法加载CA私钥");
            return;
        }
        if (entityCert == null) {
            System.err.println("无法加载实体证书");
            return;
        }

        // 2. 创建OCSP请求
        makereq(entityCert, "ocsp.req");

        // 3. 生成OCSP响应
        makeres("ocsp.req", "ocsp.res", caCert, caPrivateKey);
    }

    // 创建OCSP请求并写入文件
    public static void makereq(X509Certificate cert, String reqFile) throws Exception {
        // 创建证书ID
        CertificateID cid = new JcaCertificateID(
                /*
                以下这行代码中设置ocsp.req的签名算法
                SM3请求算法:new AlgorithmIdentifier(GMObjectIdentifiers.sm3)
                SHA1请求算法:CertificateID.HASH_SHA1
                */
                new JcaDigestCalculatorProviderBuilder().build().get(new AlgorithmIdentifier(GMObjectIdentifiers.sm3)),
                cert,
                cert.getSerialNumber()
        );

        // 构建OCSP请求
        OCSPReqBuilder reqBuilder = new OCSPReqBuilder();
        reqBuilder.addRequest(cid);

        // 生成请求
        OCSPReq req = reqBuilder.build();

        // 写入文件
        try (FileOutputStream fos = new FileOutputStream(reqFile)) {
            fos.write(req.getEncoded());
        }
        System.out.println("OCSP请求已写入: " + reqFile);
    }

    // 生成OCSP响应并写入文件
    public static void makeres(String reqFile, String resFile, 
                               X509Certificate caCert, PrivateKey caKey) 
            throws Exception {
        
        // 1. 读取请求文件
        byte[] reqData = readFile(reqFile);
        OCSPReq req = new OCSPReq(reqData);

        // 2. 准备响应生成器
        RespID respID = new RespID(new X500Name(caCert.getSubjectX500Principal().getName()));
        BasicOCSPRespBuilder respBuilder = new BasicOCSPRespBuilder(respID);

        // 3. 处理每个请求
        Req[] requests = req.getRequestList();
        for (Req ocspReq : requests) {
            CertificateID certId = ocspReq.getCertID();
            BigInteger serial = certId.getSerialNumber();

            // 创建证书状态 (Good)
            CertificateStatus status = CertificateStatus.GOOD;
            
            // 添加响应信息 (当前时间，有效期1天)
            respBuilder.addResponse(certId, status, new Date(), new Date(System.currentTimeMillis() + 86400000), null);
        }

        // 4. 创建内容签名器 - 根据CA证书类型选择签名算法
        String signatureAlgorithm = getSignatureAlgorithm(caCert);
        ContentSigner signer = new JcaContentSignerBuilder(signatureAlgorithm)
                .setProvider("BC")
                .build(caKey);

        // 5. 生成响应
        BasicOCSPResp basicResp = respBuilder.build(signer, 
                new X509CertificateHolder[]{new X509CertificateHolder(caCert.getEncoded())}, 
                new Date());

        // 6. 包装成完整响应
        OCSPRespBuilder ocspRespBuilder = new OCSPRespBuilder();
        OCSPResp resp = ocspRespBuilder.build(OCSPRespBuilder.SUCCESSFUL, basicResp);

        // 7. 写入文件
        try (FileOutputStream fos = new FileOutputStream(resFile)) {
            fos.write(resp.getEncoded());
        }
        System.out.println("OCSP响应已写入: " + resFile);
    }

    // 辅助方法：读取文件内容
    private static byte[] readFile(String filename) throws IOException {
        try (FileInputStream fis = new FileInputStream(filename);
             ByteArrayOutputStream bos = new ByteArrayOutputStream()) {
            byte[] buffer = new byte[1024];
            int len;
            while ((len = fis.read(buffer)) > 0) {
                bos.write(buffer, 0, len);
            }
            return bos.toByteArray();
        }
    }

    // 加载证书 (支持PEM和DER格式)
    private static X509Certificate loadCert(String filename) throws Exception {
        try (FileInputStream fis = new FileInputStream(filename)) {
            // 尝试PEM格式
            if (isPEMFile(filename)) {
                PEMParser parser = new PEMParser(new InputStreamReader(fis));
                Object obj = parser.readObject();
                if (obj instanceof X509CertificateHolder) {
                    return new JcaX509CertificateConverter()
                            .setProvider("BC")  // 使用BC提供程序
                            .getCertificate((X509CertificateHolder) obj);
                }
                throw new IOException("不支持的PEM格式: " + obj.getClass().getName());
            }
            
            // 尝试DER格式
            CertificateFactory cf = CertificateFactory.getInstance("X.509", "BC"); // 使用BC提供程序
            return (X509Certificate) cf.generateCertificate(fis);
        }
    }

    // 加载私钥 (支持PEM格式)
    private static PrivateKey loadKey(String filename) throws Exception {
        try (FileReader reader = new FileReader(filename)) {
            PEMParser parser = new PEMParser(reader);
            Object obj = parser.readObject();
            JcaPEMKeyConverter converter = new JcaPEMKeyConverter().setProvider("BC");
            
            if (obj instanceof org.bouncycastle.openssl.PEMKeyPair) {
                // 处理传统RSA私钥 (PKCS#1格式)
                return converter.getPrivateKey(((org.bouncycastle.openssl.PEMKeyPair) obj).getPrivateKeyInfo());
            } else if (obj instanceof org.bouncycastle.asn1.pkcs.PrivateKeyInfo) {
                // 处理PKCS#8格式私钥
                return converter.getPrivateKey((org.bouncycastle.asn1.pkcs.PrivateKeyInfo) obj);
            } else if (obj instanceof org.bouncycastle.openssl.PEMEncryptedKeyPair) {
                throw new IOException("加密私钥需要密码");
            }
            
            // 尝试直接作为私钥对象处理
            if (obj instanceof PrivateKey) {
                return (PrivateKey) obj;
            }
            
            throw new IOException("不支持的私钥格式: " + (obj != null ? obj.getClass().getName() : "null"));
        }
    }

    // 检查文件是否为PEM格式
    private static boolean isPEMFile(String filename) {
        try (BufferedReader reader = new BufferedReader(new FileReader(filename))) {
            String firstLine = reader.readLine();
            return firstLine != null && firstLine.contains("-----BEGIN");
        } catch (IOException e) {
            return false;
        }
    }
    
    // 根据证书类型确定签名算法
    private static String getSignatureAlgorithm(X509Certificate cert) {
        String algorithm = cert.getSigAlgName();
        if (algorithm.contains("SM3")) {
            return "SM3withSM2";
        } else if (algorithm.contains("ECDSA")) {
            return "SHA256withECDSA";
        } else {
            return "SHA256withRSA";
        }
    }
}